from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings

# 创建向量存储
embeddings = OpenAIEmbeddings(
    api_key="sk-VowKQBUMIkSND8WScNJtDLqf3FyqWHQ43LMVUXH1m6GZaopA",
    base_url="https://ai.nengyongai.cn/v1",
    model="text-embedding-3-small"
)
#
texts = [
    "笨笨是一只很喜欢睡觉的猫咪",
    "我喜欢在夜晚听音乐，这让我感到放松。",
    "猫咪在窗台上打盹，看起来非常可爱。",
    "学习新技能是每个人都应该追求的目标。",
    "我最喜欢的食物是意大利面，尤其是番茄酱的那种。",
    "昨晚我做了一个奇怪的梦，梦见自己在太空飞行。",
    "我的手机突然关机了，让我有些焦虑。",
    "阅读是我每天都会做的事情，我觉得很充实。",
    "他们一起计划了一次周末的野餐，希望天气能好。",
    "我的狗喜欢追逐球，看起来非常开心。",
]
#
vector_store = FAISS.from_texts(texts, embeddings)


def test1():
    print(vector_store.index.ntotal)


def test2():
    """
    基础相似度搜索，传递 query(搜索语句)、k(返回条数)、filter(过滤器)、fetch_k(富余条数) 等
    """
    print(vector_store.similarity_search("我养了一只猫，叫笨笨"))


def test3():
    """
    携带得分的相似性搜索，参数和 similarity_search() 函数保持一致，
    只是会返回得分，这里的得分并不是相似性得分，而是欧几里得距离。
    """
    print(vector_store.similarity_search_with_score("我养了一只猫，叫笨笨"))


def test4():
    """
    携带相关性得分的相似性搜索，得分范围是 0-1
    """
    print(vector_store.similarity_search_with_relevance_scores("我养了一只猫，叫笨笨"))
